import torch
from torch.autograd import Variable
import os
import numpy as np
from sklearn.model_selection import train_test_split
import cv2
from torch import nn

torch.manual_seed(777)  # reproducibility

DATADIR = "C:\DL_DATA\dl2_DataSets\data3"

def readData(path):
    x_images = []
    y_labels = []
    for image_name in os.listdir(path):
        image_path = os.path.join(path, image_name)
        image = cv2.imread(image_path) / 255.
        x_images.append(image)
        y_labels.append(int(image_name[0]))
    return np.array(x_images), np.array(y_labels)

x_images, y_labels = readData(DATADIR)

x_images = Variable(torch.Tensor(x_images))
y_labels = Variable(torch.Tensor(y_labels))
x_images = torch.transpose(x_images,1,3)

x_train, x_test, y_train, y_test = train_test_split(x_images, y_labels, test_size=0.3)

class conv2dCell(nn.Module):
    def __init__(self, in_ch, out_ch, kernel, stride=1, padding=0):
        super(conv2dCell, self).__init__()
        self.conv2d = nn.Sequential(
            nn.Conv2d(in_ch, out_ch, kernel, stride, padding),
            nn.BatchNorm2d(out_ch, eps=1e-3),
            nn.ReLU(True)
        )
    def forward(self, x):
        x = self.conv2d(x)
        return x

class inceptionBlock(nn.Module):
    def __init__(self, in_ch, out_ch, stride=1):
        super(inceptionBlock, self).__init__()
        self.branch1x1 = conv2dCell(in_ch, out_ch, 1, stride=stride)
        self.branch3x3 = nn.Sequential(
            conv2dCell(in_ch, out_ch, 1, stride=stride),
            conv2dCell(out_ch, out_ch, 3, padding=1)
        )
        self.branch5x5 = nn.Sequential(
            conv2dCell(in_ch, out_ch, 1, stride=stride),
            conv2dCell(out_ch, out_ch, 5, padding=2)
        )
        self.branch_pool = nn.Sequential(
            nn.MaxPool2d(3, stride=stride, padding=1),
            conv2dCell(in_ch, out_ch, 1)
        )
    def forward(self, x):
        f1 = self.branch1x1(x)
        f2 = self.branch3x3(x)
        f3 = self.branch5x5(x)
        f4 = self.branch_pool(x)
        output = torch.cat((f1, f2, f3, f4), dim=1)
        return output

class InceptionNet10(nn.Module):
    def __init__(self, in_channel, num_classes):
        super(InceptionNet10, self).__init__()
        self.conv = nn.Sequential(
            conv2dCell(in_channel, 16, kernel=3, padding=1),
            nn.MaxPool2d(1, 1)
        )
        self.blocks = nn.Sequential(
            inceptionBlock(16, 32, stride = 2),
            inceptionBlock(128, 32),

            inceptionBlock(128, 64, stride = 2),
            inceptionBlock(256, 64)
        )

        self.avgpool = nn.AvgPool2d(kernel_size=8, stride=1)
        self.classifier = nn.Linear(256, num_classes)

    def forward(self, x):
        x = self.conv(x)
        x = self.blocks(x)
        x = self.avgpool(x)
        x = torch.squeeze(x)
        x = self.classifier(x)
        return x


if __name__ == '__main__':

    model = InceptionNet10(in_channel=3, num_classes=3)

    criterion=torch.nn.CrossEntropyLoss()
    optimizer=torch.optim.Adam(model.parameters(), lr=0.001)

    for i in range(10):
        optimizer.zero_grad()
        h = model(x_train)
        cost = criterion(h, y_train.long())
        cost.backward()
        optimizer.step()

        print(i + 1, cost.item())

    pred = torch.argmax(model(x_train), 1)
    acc = (y_train == pred).float().mean()
    print('acc', acc.item())

    pred = torch.argmax(model(x_test), 1)
    acc = (y_test == pred).float().mean()
    print('acc', acc.item())
